#include "StdAfx.h"
#include "WMIEventSink.h"
#include "WMIHandler.h"

CWMIEventSink::CWMIEventSink(void)
: m_lRef(0)
, m_pUnsecApp(NULL)
, m_pStubSink(NULL)
, m_bRegistered(false)
, m_pWbemSvc(NULL)
{
  AddRef();
}

CWMIEventSink::~CWMIEventSink(void)
{
  if (m_lRef != 0)
  {
    return;
  }

  if (isRegistered())
  {
    unregisterSelf();
  }
}

ULONG CWMIEventSink::AddRef()
{
  return InterlockedIncrement(&m_lRef);
}

ULONG CWMIEventSink::Release()
{
  LONG lRef = InterlockedDecrement(&m_lRef);
  if(lRef == 0)
    delete this;
  return lRef;
}

HRESULT CWMIEventSink::QueryInterface(REFIID riid, void** ppv)
{
  if (riid == IID_IUnknown || riid == IID_IWbemObjectSink)
  {
    *ppv = (IWbemObjectSink *) this;
    AddRef();
    return WBEM_S_NO_ERROR;
  }
  else return E_NOINTERFACE;
}

ULONG CWMIEventSink::registerSelf(LPCTSTR lpcszQuery, IWbemServices *pWbemSvc)
{
  if ( isRegistered() )
  {
    return WMI_HANDLER_RTCODE_EVENTSINK_HAS_BEEN_REGISTERED;
  }

  HRESULT hRes;

  m_pWbemSvc = pWbemSvc;
  m_pWbemSvc->AddRef();

  hRes = CoCreateInstance(
    CLSID_UnsecuredApartment,
    NULL,
    CLSCTX_LOCAL_SERVER,
    IID_IUnsecuredApartment,
    (void **)&m_pUnsecApp );
  if ( FAILED(hRes) )
  {
    m_pWbemSvc->Release();
    m_pWbemSvc = NULL;

    return WMI_HANDLER_RTCODE_CREATE_UNSECUREDAPARTMENT_FAILED;
  }

  IUnknown *pStubUnk = NULL;
  m_pUnsecApp->CreateObjectStub(this, &pStubUnk);

  pStubUnk->QueryInterface(IID_IWbemObjectSink, (void **)&m_pStubSink);
  pStubUnk->Release();

  hRes = m_pWbemSvc->ExecNotificationQueryAsync(
    TEXT("WQL"),
    _bstr_t(lpcszQuery),
    WBEM_FLAG_SEND_STATUS,
    NULL,
    m_pStubSink);
  if ( FAILED(hRes) )
  {
    m_pStubSink->Release();
    m_pStubSink = NULL;

    m_pUnsecApp->Release();
    m_pUnsecApp = NULL;

    m_pWbemSvc->Release();
    m_pWbemSvc = NULL;

    return WMI_HANDLER_RTCODE_EXEC_NOTIFYQUERY_FAILED;
  }

  m_bRegistered = true;
  
  return WMI_HANDLER_RTCODE_OK;
}

ULONG CWMIEventSink::unregisterSelf(void)
{
  if ( !isRegistered() )
  {
    return WMI_HANDLER_RTCODE_OK;
  }
  
  HRESULT hRes;

  hRes = m_pWbemSvc->CancelAsyncCall(m_pStubSink);
  
  m_pStubSink->Release();
  m_pStubSink = NULL;

  m_pUnsecApp->Release();
  m_pUnsecApp = NULL;

  m_pWbemSvc->Release();
  m_pWbemSvc = NULL;

  m_bRegistered = false;

  return WMI_HANDLER_RTCODE_OK;
}
